[Draft][MoE] Expert pool decode path: validation in the step kernel, one per-forward assertion, fewer launches - #50

Scope: fork-side development on 01554/vllm; not part of upstream PR vllm-project#56177. Whether any of this goes upstream is a separate, later decision. Measurements in this PR (profiles, fresh3) are evidence for the P1 integration head 205983f only and are not attributed to the vllm-project#56177 head.

Draft for team review (author claude, review 1 mid, review 2 + GPU astra). Base: lab/upstream-pool-standalone-rebased (5119ca4, the vllm-project#56177 head). Independent of Draft #49 (tables.py staged-suffix vectorization); not to be mixed. Five commits: validation (9e1e8fe), remap kernel (60c80ee), tests (1bd6c8a), test fixtures (2e2355e: per-test distributed fixture re-init, staging width, subprocess exit path; e65bcc4: exit from inside the handler); production code is in the first two so they can be measured separately. A GPU regression on 1bd6c8a had 43 passed / 11 failed, all fixture-caused (module-scoped distributed fixture torn down per test, CPU staging width, SIGABRT during unwinding after the device assertion); no production failure.

Purpose

E6 decode profile (results/goal66-decode-profile-e6-window16-20260910, lab note e6-decode-profile-analysis-2026-09-10.md): 3,684 GPU events per token (kernel + memcpy + memset trace events, of which 3,658 are kernels; not host launch calls), of which the kernels matching the pool's per-layer route-validation chain are about 1,014 per token (1.28 ms of kernel time as a trace category sum, not a wall-clock saving). This PR keeps validation on every request and moves it into the device planner, removes the route mask and the torch remap chain on the decode path, and keeps the wide path unchanged.

Commit 1: validation in the step kernel, sanitized routes, one assertion per call

• Every lane validated in the step kernel (tables.py): out-of-range ids were already masked before any table read or write; router weights are now validated too (finite, nonnegative). A lane failing either check is planned as padding: no table read, no ownership mutation, and the placement after the step equals the same step with that lane as -1. Duplicate routes stay legal. The sticky error is set and mirrored in a device ok flag.
• Sanitized routes for every consumer (safe_ids, written by the kernel): the decode consumer passes safe_ids (invalid lanes as -1) to align, Marlin and the activation, so no raw id reaches them (this closes the activation-path expert_map[expert_id] read that only masks >= 0). Because every valid route is present in the step map after the step (hit, promoted, or staged; staging capacity equals the plan width), the route mask is the identity on decode and is skipped.
• One device assertion per call (layer.py): torch._assert_async(tables.ok) right after the step, one launch instead of about eleven per layer. Same detection path as the previous per-layer _assert_async (device assertion surfaces at the caller's next synchronization), and it holds for full forwards, single-layer and partial executions alike. clear_error(tables) resets the flag after handling.
• Wide (partition) path: unchanged, keeps its per-call assertion chain and mask_routes.

Commit 2: physical_block_experts_device

The torch chain (arange, compare, where, clamp, index, where) becomes one Triton launch on CUDA; identical semantics including blocks beyond post_padded (never indexed, -1) and absent experts (-1). The torch version remains the reference and the CPU path.

Tests

• test_pool_tables.py (CPU reference): invalid ids leave the tables identical to the same step with those lanes as padding (gate closed and open), route the other lanes, hide the lane in safe_ids, set error and clear ok; NaN/±Inf/negative weights do the same, the same weight on a padding lane is never loaded, duplicates with finite weights never set the error and resolve to one row; 40 random trials (gate, staging shortage, duplicates, sentinels, invalid ids and weights) assert every valid lane is routed in routes, safe_ids and the step map, and clear_error restores a consistent pool.
• test_pool_layer_helpers.py: kernel-vs-torch equality of the remap on CUDA (bank > E, absent experts, garbage ids beyond post_padded), CPU fallback.
• test_pool_marlin_cuda.py: (a) step-only device test, no consumer so no assertion fires: out-of-range id, negative non-sentinel id and NaN weight at layers 0/1/2 leave the tables equal to twin pools stepped with padding, set error/ok, stay sticky across clean steps on other layers until clear_error; (b) the decode path captured in a CUDA graph and replayed with four different valid inputs matches eager results on identically placed twin pools (only the two warm-up steps are mirrored; placement, step map and safe_ids are compared after every replay); (c) the planner step alone captured in a graph: an invalid replay sets the sticky error, a clean replay keeps it, clear_error resets it, placement matches a padding twin throughout; (d) with a clamp-limited activation config, the consumer called directly on the kernel's safe_ids after out-of-range / negative / all-invalid steps produces exactly the padding oracle's output (no raw id reaches the activation's expert_map read); (e) five subprocess-isolated assertion cases: invalid id at the first layer, NaN weight at a middle layer, negative id at the last layer, a single non-final layer applied alone then synchronized (partial execution), and an invalid id written into a captured graph's input buffer before replay. Clean setup, capture and clean replay run outside the guarded region; only the device-side assertion counts as the expected failure, any other error fails the case. The existing two-layer decode/prefill equivalence test is unchanged.

Local status: pre-commit (ruff, format, mypy 3.10, SPDX, torch.cuda guard) passed on all commits. clear_error resets the flags only; a fired device assertion poisons the CUDA context and needs a process restart (documented). The test files could not be executed on the author's machine (import torch is killed there, exit 137). GPU execution on e65bcc4 (astra, results/expert-pool-validation-e65bcc45d-20260910): 54 passed / 0 skipped plus bench unittest 2, all five subprocess cases recorded the expected device assertion with rc 0, final exit 0 / not OOM-killed.

Profile results so far (diagnostic runs, same recipe as E6; lab note e6-decode-profile-analysis-2026-09-10.md)

Per token, average of the same 14-step window, kernel-duration sums by name category (not wall-clock attribution):

The copy category differs strongly between runs (17 % / 29 % / 59 % of copy launches above 10 µs in the same 14-step window; a duration-threshold count, transfer counts are not recorded), following each run's generation trajectory (outputs differ). The window wall (14.81 / 15.65 / 20.73 ms) shows no wall reduction in the observed windows, and because trajectories and copy activity differ between runs the causal effect of the change cannot be separated from the wall; category sums overlap across streams and are not attributed to wall. The launch reduction and the disappearance of the validation chain are visible in every run. Speed is decided by a non-profiled fresh3 on 205983f against the accepted 63.19 tok/s under identical conditions (result below).

fresh3 result (205983f, results/goal66-p1-remap-repeat3-20260910)

Acceptance recipe, identical conditions: frozen pair 3d4e64ec, N258, max-model-len 4096, three fresh server starts, warmup A then measure B, temperature 0. Every run: exit 0, not OOM-killed, both requests finished with stop, no errors; provenance candidate 205983f on base a97dacb. Executed by astra (root recomputation), independently verified by the author from the originals.

Median 69.10 tok/s versus the accepted 63.19 (7dedc6d, median of 63.2 / 62.7 / 63.7 under the same recipe): +9.4 %. This is a historical comparison across separate starts with differing output trajectories, so it is evidence for the integrated head 205983f as a whole and is not attributed to the validation or remap commit individually. The profile findings above (−42 % GPU events, no wall improvement observed in the windows) stand unchanged.

Review outcome (2026-09-10): review 1 (mid) and review 2 (astra) both accept 205983f as a short-decode speed improvement candidate on the fork, adopting the median 69.10 tok/s of the whole-request observed decode. Limits kept by both: the +9.4 % is a historical comparison on a different trajectory, not attributed to an individual commit, to a same-trajectory speedup, or to the current upstream vllm-project#56177 head; this is not a re-acceptance of the SWELancer / long-output / long-context quality checks on the new head, not an upstream adoption decision, and the FreeToken gap is not closed.

Measurement plan (astra, completed)

Same recipe as the E6 profile (N258, ctx 4096, frozen pair): profile run per commit for the event count and category changes (tables above), then fresh3 against the accepted 63.19 under identical conditions (result above). Both reviews recorded above.

🤖 Generated with Claude Code

https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT